const os = require('os');
const fs = require('fs');
const path = require('path');
const sqlite3 = require('sqlite3').verbose();
const axios = require('axios');
const yargs = require('yargs/yargs');
const { hideBin } = require('yargs/helpers');
const archiver = require('archiver');
const sdk = require('@baiducloud/sdk');
const BosClient = sdk.BosClient;
const { DatasetsTable } = require('../../main/dbs/tables/tableDatasets');
const { CategoriesTable } = require('../../main/dbs/tables/tableCategories');
const { DatasetsTableCloud } = require('../../main/dbs/tables/tableDatasetsCloud');
const { AutoUpdater } = require('../../main/utils/autoupdate');

class DatasetsTool {
    constructor() {
        // bos相关的属性
        this.bos_config = {
            endpoint: 'http://bd.bcebos.com',
            credentials: {
                ak: '6356708ce8d342bf965f2047d7c3cdb2',
                sk: 'd4837f696fdf4cf4abf11a80a1067abb'
            }
        };
        this.bos_client = new BosClient(this.bos_config);
        this.bos_bucket = 'label-anything';

        // 数据库相关属性
        let userHomeDir = os.homedir();
        let appName = 'labelapp';
        this.userDataPath = null;
        switch (process.platform) {
            case 'darwin': // macOS
                this.userDataPath = path.join(userHomeDir, 'Library', 'Application Support', appName);
                break;
            case 'win32': // Windows
                this.userDataPath = path.join(process.env.APPDATA, appName);
                break;
            default: // Linux and others
                this.userDataPath = path.join(userHomeDir, `.config`, `${appName}`);
        }
    }

    async listDatasetsCloud() {
        let datasets_cloud = await this.dataset_cloud_table.getDatasetList();
        console.log(datasets_cloud);
        return datasets_cloud;
    }

    // 从本地的clouddb删除记录
    async deleteFromCloudDB(dataset_id) {
        await this.dataset_cloud_table.deleteDataset(dataset_id);
    }

    async listDatasetsLocal() {
        let dataset_list = this.dataset_table && (await this.dataset_table.getDatasetList());
        console.log(dataset_list);
        return dataset_list;
    }

    async uploadDataset(dataset_id, desc) {
        // 获取dataset信息
        console.log(`## Upload dataset '${dataset_id}' with desc: ${desc}`);
        let dataset_info;
        try {
            dataset_info = await this.dataset_table.getDatasetInfo(dataset_id);
        } catch (error) {
            console.log(`## Dataset ${dataset_id} not found`);
            return;
        }

        // 更新数据集信息到云端数据集sqlite
        let localDir = dataset_info.dataset_dir;
        let remoteDir = `/labelapp/datasets_cloud/${dataset_id}`;
        if (!(await this.updateCloudVersion(dataset_info, desc))) {
            return;
        }

        // 上传数据集类别信息
        await this.saveCategoriesInfo(dataset_info);

        // 上传数据集目录，包括list.db、categroies.json
        let localDBPath = path.join(this.userDataPath, 'datasets', `${dataset_id}`, 'list.db');
        fs.copyFileSync(localDBPath, path.join(localDir, 'list.db'));
        // await this.uploadDirectory(localDir, remoteDir, ['.DS_Store']);

        // 压缩并上传数据集
        const zipFilePath = path.join(this.userDataPath, 'datasets', `${dataset_id}.zip`);
        await this.zipDirectory(localDir, zipFilePath, ['.DS_Store']);
        await this.uploadFile(zipFilePath, `/labelapp/datasets_cloud/${dataset_id}/${dataset_id}.zip`);
    }

    // 保存数据集的类别信息
    async saveCategoriesInfo(dataset_info) {
        let dataset_id = dataset_info.id;
        let localDir = dataset_info.dataset_dir;

        let categories = await this.categories_table.getCategories(dataset_id);
        let categories_info = [];
        for (let i = 0; i < categories.length; i++) {
            let category = categories[i];
            let category_info = {
                category_id: category.category_id,
                category_name: category.category_name,
                supercategory: category.supercategory,
                category_name: category.category_name
            };
            categories_info.push(category_info);
        }

        try {
            let categoriesInfoPath = path.join(localDir, 'categories.json');
            fs.writeFileSync(categoriesInfoPath, JSON.stringify(categories_info));
            return true;
        } catch (error) {
            console.error('saveCategoriesInfo error = ', error);
            return false;
        }
    }

    async updateCloudVersion(localDatasetInfos, datasetDesc) {
        // console.log(`localDatasetInfos = `, localDatasetInfos);
        // 更新数据库记录
        try {
            if (localDatasetInfos) {
                let dataset_id = localDatasetInfos.id;
                let downloadUrl = `https://label-anything.bd.bcebos.com/labelapp/datasets_cloud/${dataset_id}/${dataset_id}.zip`;
                let datasetInfos = {
                    ...localDatasetInfos
                };
                datasetInfos.dataset_dir = `/labelapp/datasets_cloud/${dataset_id}`;
                datasetInfos.download_url = downloadUrl;
                datasetInfos.dataset_desc = datasetDesc;

                // console.log(datasetInfos);
                await this.dataset_cloud_table.deleteDataset(dataset_id);
                await this.dataset_cloud_table.insertDataset(datasetInfos);
            }
        } catch (error) {
            console.error('updateCloudVersion error = ', error);
        }

        // 更新版本号
        let currentVersion = await this.getCurrentVersion();
        if (!currentVersion) {
            console.error(`Error getting current version from BOS!`);
            return false;
        }
        let newVersion = this.incrementVersion(currentVersion);
        let newVersionPath = path.join(this.userDataPath, 'datasets_cloud', 'VERSION');
        fs.writeFileSync(newVersionPath, newVersion);

        // 压缩datasets_cloud压缩包
        const packDir = path.join(this.userDataPath, 'datasets_cloud');
        const zipFilePath = path.join(this.userDataPath, `datasets-${newVersion}.zip`);
        await this.zipDirectory(packDir, zipFilePath, ['.DS_Store']);
        console.log(`## Directory compressed to ${zipFilePath}`);

        // 上传压缩文件到百度云BOS
        let remotePath = `/labelapp/datasets_cloud/datasets-${newVersion}.zip`;
        await this.uploadFile(zipFilePath, remotePath);
        console.log(`## File ${remotePath} uploaded to BOS.`);

        // 上传新的版本号到VERSION
        await this.uploadFile(newVersionPath, '/labelapp/datasets_cloud/VERSION');
        console.log(`## Updated version info to BOS.`);

        // 清理临时压缩文件
        // fs.unlinkSync(zipFilePath);
        console.log('## Temporary file cleaned up.');

        return true;
    }

    async getCurrentVersion() {
        try {
            let index_url = 'https://label-anything.bd.bcebos.com/labelapp/datasets_cloud/VERSION';
            let resp = await axios.get(index_url);
            return resp.data.trim();
        } catch (error) {
            return null;
        }
    }

    incrementVersion(version) {
        const parts = version.split('.').map(Number);
        parts[2] += 1;
        return parts.join('.');
    }

    async uploadDirectory(localDir, remoteDir, blacklist = []) {
        const files = fs.readdirSync(localDir, { withFileTypes: true });

        for (const file of files) {
            if (blacklist.includes(file.name)) {
                // console.log(`Skipping blacklisted ${file.name}`);
                continue;
            }

            const localFilePath = path.join(localDir, file.name);
            const remoteFilePath = `${remoteDir}/${file.name}`;

            if (file.isDirectory()) {
                await this.uploadDirectory(localFilePath, remoteFilePath, blacklist);
            } else {
                await this.uploadFile(localFilePath, remoteFilePath);
            }
        }
    }

    async uploadFile(localFilePath, remoteFilePath) {
        console.log(`## Uploading file to ${remoteFilePath}`);
        await this.uploadFileToBOS(localFilePath, remoteFilePath);
    }

    async initDB() {
        // 更新云端数据集
        console.log('##====================================');
        const updater = new AutoUpdater({
            userDataPath: this.userDataPath
        });
        await updater.update('datasets_cloud');
        console.log('##------------------------------------');

        // 初始化datasets_cloud数据库
        const dbCloudPath = path.join(this.userDataPath, 'datasets_cloud', 'datasets_cloud.db');
        this.db_cloud = new sqlite3.Database(dbCloudPath);
        console.log(`## Open DatasetsCloudDB OK [${dbCloudPath}]`);

        // 初始化DatasetsTableCloud表
        this.dataset_cloud_table = new DatasetsTableCloud(null, this.db_cloud);
        this.dataset_cloud_table.init();

        // 初始化labelstudio.db数据库
        const dbFilePath = path.join(this.userDataPath, 'labelstudio.db');
        this.db = new sqlite3.Database(dbFilePath);
        console.log(`## Open LabelStudioDB OK [${dbFilePath}]`);

        // 初始化DatasetsTable表
        this.dataset_table = new DatasetsTable(null, this.db);
        this.dataset_table.init();

        // 初始化CategoriesTable表
        this.categories_table = new CategoriesTable(null, this.db);
        this.categories_table.init();
        console.log('##------------------------------------');
    }

    async close() {
        console.log('##------------------------------------');
        this.db && (await this.closeDB());
        this.db_cloud && (await this.closeDBCloud());
        console.log('##====================================');
    }

    async closeDB() {
        return new Promise((resolve, reject) => {
            if (!this.db) {
                reject(new Error('Database is not initialized.'));
                return;
            }

            this.db.close((err) => {
                if (err) {
                    console.error('## Close LabelStudioDB Failed, Error: ', err);
                    reject(err);
                } else {
                    console.log('## Close LabelStudioDB OK!');
                    resolve();
                }
            });
        });
    }

    async closeDBCloud() {
        return new Promise((resolve, reject) => {
            if (!this.db_cloud) {
                reject(new Error('Database is not initialized.'));
                return;
            }

            this.db_cloud.close((err) => {
                if (err) {
                    console.error('## Close DatasetsCloudDB Failed, Error: ', err);
                    reject(err);
                } else {
                    console.log('## Close DatasetsCloudDB OK!');
                    resolve();
                }
            });
        });
    }

    // BOS相关操作
    async uploadFileToBOS(localFilePath, remoteFileName) {
        try {
            const response = await this.bos_client.putObjectFromFile(this.bos_bucket, remoteFileName, localFilePath);
            return response;
        } catch (error) {
            console.error(`Failed to upload ${remoteFileName}:`, error);
            throw error;
        }
    }

    // 压缩文件夹
    zipDirectory(sourceDir, outPath, blacklist = []) {
        const archive = archiver('zip', { zlib: { level: 9 } });
        const stream = fs.createWriteStream(outPath);

        return new Promise((resolve, reject) => {
            archive.on('error', (err) => reject(err));

            // 使用递归函数来添加目录并检查黑名单
            const addDirectory = (dir, base) => {
                fs.readdirSync(dir).forEach((file) => {
                    const filePath = `${dir}/${file}`;
                    const relativePath = base ? `${base}/${file}` : file;

                    if (blacklist.includes(file)) {
                        return; // 如果在黑名单中，跳过
                    }

                    if (fs.lstatSync(filePath).isDirectory()) {
                        addDirectory(filePath, relativePath);
                    } else {
                        archive.file(filePath, { name: relativePath });
                    }
                });
            };

            addDirectory(sourceDir, '');

            archive.pipe(stream);
            stream.on('close', () => resolve());
            archive.finalize();
        });
    }
}

(async () => {
    const argv = yargs(hideBin(process.argv))
        .command('list_local', 'List all local datasets')
        .command('list_cloud', 'List all cloud datasets')
        .command('del_cloud', 'delete record from clouddb', {
            id: {
                description: 'The ID of the dataset',
                alias: 'i',
                type: 'string',
                demandOption: true
            }
        })
        .command('add_cloud', 'add local dataset to clouddb', {
            id: {
                description: 'The ID of the dataset',
                alias: 'i',
                type: 'string',
                demandOption: true
            },
            desc: {
                description: 'Description for the local dataset',
                alias: 'd',
                type: 'string',
                demandOption: true
            }
        })
        .command('update', 'Just update cloud datasets_cloud.db/version files')
        .demandCommand(1, 'You need at least one command before moving on')
        .help()
        .alias('help', 'h').argv;

    const dataset_tool = new DatasetsTool();
    await dataset_tool.initDB();

    // 根据命令执行对应的逻辑
    try {
        if (argv._.includes('list_local')) {
            await dataset_tool.listDatasetsLocal();
        } else if (argv._.includes('list_cloud')) {
            await dataset_tool.listDatasetsCloud();
        } else if (argv._.includes('del_cloud')) {
            await dataset_tool.deleteFromCloudDB(argv.id);
        } else if (argv._.includes('add_cloud')) {
            await dataset_tool.uploadDataset(argv.id, argv.desc);
        } else if (argv._.includes('update')) {
            // await dataset_tool.updateCloudVersion();
        }
    } catch (err) {
        console.error('Error:', err);
    } finally {
        await dataset_tool.close();
    }
})();
